package com.sheepone.util.structure.tree.binary;

import com.sheepone.util.enums.CompareType;
import com.sheepone.util.structure.node.AVLTreeNode;

/**
 * @author Miss.杨
 * @description AVL树（二叉搜索平衡树）
 * @since 2024/5/9 - 15:22
 */
public class AVLTree<K extends Comparable<K>, V> extends BinarySearchTree<K, V> {

    @Override
    public void put(K key, V value) {
        root = doPut((AVLTreeNode<K, V>) root, key, value);
    }

    @Override
    public void remove(K key) {
        root = doRemove((AVLTreeNode<K, V>) root, key);
    }

    /**
     * 在以root节点作为根节点的树中删除key对应的节点，
     * 返回删除后的该树的新根节点
     *
     * @param root 当前树的根节点
     * @param key  待删除节点的key值
     * @return 删除后的该树的新根节点
     */
    private AVLTreeNode<K, V> doRemove(AVLTreeNode<K, V> root, K key) {
        if (root == null) {
            return null;
        }
        CompareType compareType = compare(key, root.getKey());
        if (compareType == CompareType.LESS_THAN) {
            root.setLeft(doRemove((AVLTreeNode<K, V>) root.getLeft(), key));
        } else if (compareType == CompareType.GREATER_THAN) {
            root.setRight(doRemove((AVLTreeNode<K, V>) root.getRight(), key));
        } else {
            // 找到key对应的节点
            AVLTreeNode<K, V> left = (AVLTreeNode<K, V>) root.getLeft();
            AVLTreeNode<K, V> right = (AVLTreeNode<K, V>) root.getRight();
            if (left == null && right == null) {
                // 待删除节点是叶子节点,直接返回，不用更新该节点高度
                return null;
            } else if (right == null) {
                root = left;
            } else if (left == null) {
                root = right;
            } else {
                // 待删除节点的左右孩子都不为null
                // 从右子树中找到后继节点替换当前被删除的节点
                root = (AVLTreeNode<K, V>) min(right);
                // 从被删除节点的右子树中删除该后继节点
                right = doRemove(right, root.getKey());
                // 让后继节点替换当前节点
                root.setRight(right);
                root.setLeft(left);
            }
        }
        // 逐级递归更新高度并平衡
        updateHeight(root);
        return balance(root);
    }

    private AVLTreeNode<K, V> doPut(AVLTreeNode<K, V> node, K key, V value) {
        // 找到空位，执行新增操作
        if (node == null) {
            size++;
            return new AVLTreeNode<>(key, value);
        }
        CompareType compareType = compare(key, node.getKey());
        if (compareType == CompareType.EQUAL) {
            node.setValue(value);
            return node;
        }
        if (compareType == CompareType.GREATER_THAN) {
            node.setRight(doPut((AVLTreeNode<K, V>) node.getRight(), key, value));
        } else {
            node.setLeft(doPut((AVLTreeNode<K, V>) node.getLeft(), key, value));
        }
        // 1.更新高度
        updateHeight(node);
        // 2.平衡高度
        return balance(node);
    }

    /**
     * 平衡当前节点
     *
     * @param node 需要平衡的节点
     * @return
     */
    private AVLTreeNode<K, V> balance(AVLTreeNode<K, V> node) {
        BalanceFactor balanceFactor = balanceFactor(node);
        if (balanceFactor == BalanceFactor.LEFT_HEAVY) {
            AVLTreeNode<K, V> left = (AVLTreeNode<K, V>) node.getLeft();
            if (balanceFactor(left) == BalanceFactor.LEFT_HEAVY || balanceFactor(left) == BalanceFactor.BALANCED) {
                node = rightRotate(node);
            } else {
                node = leftRightRotate(node);
            }
        } else if (balanceFactor == BalanceFactor.RIGHT_HEAVY) {
            AVLTreeNode<K, V> right = (AVLTreeNode<K, V>) node.getRight();
            if (balanceFactor(right) == BalanceFactor.RIGHT_HEAVY || balanceFactor(right) == BalanceFactor.BALANCED) {
                node = leftRotate(node);
            } else {
                node = rightLeftRotate(node);
            }
        }
        return node;
    }

    private int height(AVLTreeNode<K, V> node) {
        return node == null ? 0 : node.getHeight();
    }

    private void updateHeight(AVLTreeNode<K, V> node) {
        node.setHeight(Math.max(height((AVLTreeNode<K, V>) node.getLeft()), height((AVLTreeNode<K, V>) node.getRight())) + 1);
    }

    /**
     * 左旋，降低右边高度
     * 1. 将当前节点的右孩子的左孩子作为当前节点的右孩子
     * 2. 将当前节点作为当前节点的右孩子的左孩子（当前节点比其右孩子要小）
     *
     * @param node 需要旋转的根节点
     * @return 旋转后的根节点
     */
    private AVLTreeNode<K, V> leftRotate(AVLTreeNode<K, V> node) {
        AVLTreeNode<K, V> right = (AVLTreeNode<K, V>) node.getRight();
        node.setRight(right.getLeft());
        right.setLeft(node);
        // 更新高度，更新次序不可更改
        updateHeight(node);// 旋转后变低了
        updateHeight(right);// 旋转后变高了
        return right;
    }

    /**
     * 右旋，降低左边高度
     * 1. 将当前节点的的左孩子的右孩子作为当前节点的左孩子
     * 2. 将当前节点作为当前节点的左孩子的右孩子（当前节点比其左孩子要大）
     *
     * @param node 需要旋转的根节点
     * @return 旋转后的根节点
     */
    private AVLTreeNode<K, V> rightRotate(AVLTreeNode<K, V> node) {
        AVLTreeNode<K, V> left = (AVLTreeNode<K, V>) node.getLeft();
        node.setLeft(left.getRight());
        left.setRight(node);
        // 更新高度，更新次序不可更改
        updateHeight(node);// 旋转后变低了
        updateHeight(left);// 旋转后变高了
        return left;
    }

    /**
     * 左右旋，降低左边高度
     * 先左旋当前节点的左孩子节点后右旋当前节点
     *
     * @param node 需要旋转的根节点
     * @return 旋转后的根节点
     */
    private AVLTreeNode<K, V> leftRightRotate(AVLTreeNode<K, V> node) {
        node.setLeft(leftRotate((AVLTreeNode<K, V>) node.getLeft()));
        return rightRotate(node);
    }

    /**
     * 右左旋，降低右边高度
     * 先右旋当前节点的右孩子节点后左旋当前节点
     *
     * @param node 需要旋转的根节点
     * @return 旋转后的根节点
     */
    private AVLTreeNode<K, V> rightLeftRotate(AVLTreeNode<K, V> node) {
        node.setRight(rightRotate((AVLTreeNode<K, V>) node.getRight()));
        return leftRotate(node);
    }

    private BalanceFactor balanceFactor(AVLTreeNode<K, V> node) {
        if (node != null) {
            int leftHeight = height((AVLTreeNode<K, V>) node.getLeft());
            int rightHeight = height((AVLTreeNode<K, V>) node.getRight());
            int balanceFactor = leftHeight - rightHeight;
            if (balanceFactor > 1) {
                return BalanceFactor.LEFT_HEAVY;
            } else if (balanceFactor < -1) {
                return BalanceFactor.RIGHT_HEAVY;
            }
        }
        return BalanceFactor.BALANCED;
    }

    /**
     * 平衡因子枚举类
     */
    private enum BalanceFactor {
        /**
         * 左倾
         */
        LEFT_HEAVY,
        /**
         * 右倾
         */
        RIGHT_HEAVY,
        /**
         * 平衡
         */
        BALANCED
    }
}
